from flask import Flask, request, jsonify
import torch
import numpy as np
import pandas as pd
from transformers import AutoModelForCausalLM, AutoTokenizer
from dkt_prompt import parse_server_request, generate_final_prompt
# from dkt_prompt_1 import parse_server_request, generate_final_prompt
from generate_dkt_rec import DKTInference, process_student_incorrect_with_rec, generate_rec_and_probs, recommend_weak_questions, calculate_average_probability
import json
import re
import socket

app = Flask(__name__)

# 初始化 DKT 模型
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dkt_model_path = './resources/dkt_model_final.pth'  # 替换为实际模型路径
knows_tree_path = './resources/knowledge_graph_重庆数学.json'
math_teaching_plan = './resources/教学计划知识点-2025数学.json'
kp_map_path = './resources/knowledge_point_mapping.csv'
# 配置题库 CSV 文件路径
question_bank_csv_path = './resources/question_bank_processed.csv'

kp_map_df = pd.read_csv('./resources/knowledge_point_mapping.csv')
kp2id = dict(zip(kp_map_df["knowledge_point"], kp_map_df["mapped_id"]))

num_questions = 157
dkt_model = DKTInference(
    model_path=dkt_model_path,
    num_questions=num_questions,
    hidden_dim=128,
    layer_num=2,
    dropout=0.2,
    cell_type='lstm',
    device=DEVICE
)

# 加载资源
with open(knows_tree_path, "r", encoding="utf-8") as f:
    knows_tree = json.load(f)

with open(math_teaching_plan, "r", encoding="utf-8") as f:
    teaching_plan = json.load(f)
    
question_bank_df = pd.read_csv(question_bank_csv_path)

# 工具函数：构建知识点集合
def build_all_knows(student_knows_raw, extra_knows=None):
    student_trained_knows = re.split(r'[;,，\s]+', student_knows_raw.strip())
    all_knows_with_prereq = set(student_trained_knows)
    for know in student_trained_knows:
        if know in knows_tree:
            prerequisites = knows_tree[know].get("prerequisites", [])
            all_knows_with_prereq.update(prerequisites)
    if extra_knows:
        all_knows_with_prereq.update(extra_knows)
    return all_knows_with_prereq
# 工具函数：构建知识点集合 - 预习
def build_all_knows(student_knows_raw, extra_knows=None):
    student_trained_knows = re.split(r'[;,，\s]+', student_knows_raw.strip())
    all_knows_with_prereq = set(student_trained_knows)
    for know in student_trained_knows:
        if know in knows_tree:
            prerequisites = knows_tree[know].get("prerequisites", [])
            all_knows_with_prereq.update(prerequisites)
    if extra_knows:
        all_knows_with_prereq.update(extra_knows)
    return all_knows_with_prereq

# 工具函数：获取推荐结果并打印
def get_recommendations_and_log(student_data, all_knows_with_prereq, pre=None):
    relevant_questions_df = question_bank_df[
        question_bank_df["知识点"].apply(lambda x: any(know in str(x) for know in all_knows_with_prereq))
    ]
    recommendations = recommend_weak_questions(student_data, relevant_questions_df, dkt_model, question_bank_csv_path, pre=pre, top_k=5)

    print("\n[推荐题目 Top-5]")
    for idx, rec in enumerate(recommendations, 1):
        prob_str = rec['掌握概率']
        avg_prob = calculate_average_probability(prob_str)
        print(f"{idx}. 题目ID: {rec['题目ID']}，掌握概率: {prob_str}" + 
              (f" (ave: {avg_prob:.3f})" if "；" in prob_str or "，" in prob_str else "") +
              f"，知识点: {rec['知识点ID']}，难度: {rec['难度']}")
    return recommendations

# 工具函数：提取教学计划中的知识点
def extract_teaching_plan_knows(arrangementID=None, programmeID=None, topicID=None):
    topic_names = []
    if arrangementID:
        for _, arrangements in teaching_plan.items():
            if arrangementID in arrangements:
                return [topic["topicName"] for topic in arrangements[arrangementID]]
        raise ValueError("arrangementID not found")
    elif programmeID:
        if programmeID not in teaching_plan:
            raise ValueError("programmeID not found")
        for _, topic_list in teaching_plan[programmeID].items():
            topic_names.extend([topic["topicName"] for topic in topic_list])
        return topic_names
    elif topicID:
        for programme in teaching_plan.values():
            for topic_list in programme.values():
                for topic in topic_list:
                    if topic["topicID"] == topicID:
                        # print(topic["topicName"])
                        return [topic["topicName"]]
        raise ValueError("topicID not found")
        
    return []

# 工具函数：获取知识点映射并组装成新的知识点id
# todo
def map_knowledge_points(student_data, kp2id):
    raw_kp_str = student_data.get("knowledge_points", "")
    sequences = raw_kp_str.strip().split(";")  # 每题的知识点序列
    mapped_sequences = []

    for seq in sequences:
        # 对每题中的多个知识点（可能逗号分隔）进行映射
        kps = [kp.strip() for kp in re.split(r'[，,]', seq) if kp.strip()]
        mapped_ids = []
        for kp in kps:
            if kp in kp2id:
                mapped_ids.append(str(kp2id[kp]))
            else:
                print(f"[警告] 找不到知识点映射: {kp}")
        mapped_sequences.append(",".join(mapped_ids))

    # 合并为与原始结构匹配的“题目用;隔开，每题内多个ID用,隔开”的结构
    student_data["mapped_knowledge_points"] = ";".join(mapped_sequences)
    return student_data


# 不存在教学计划时调用的方法
@app.route("/recommend_weak_questions", methods=["POST"])
def recommend_weak_questions_endpoint():
    try:
        student_data = request.json
        student_data = map_knowledge_points(student_data, kp2id)
        # print("student_data",student_data)
        all_knows = build_all_knows(student_data["knowledge_points"])
        top5 = get_recommendations_and_log(student_data, all_knows)
        if top5:
            return jsonify({"status": "success", "recommendations": top5})
        return jsonify({"status": "error", "message": "没有推荐的题目"}), 400
    except Exception as e:
        return jsonify({"status": "error", "message": str(e)}), 500

# 教学计划中调用的方法
@app.route("/recommend_weak_questions/teaching/arrangement/<arrangementID>", methods=["POST"])
def recommend_weak_questions_in_teaching_plan_endpoint(arrangementID):
    try:
        student_data = request.json
        student_data = map_knowledge_points(student_data, kp2id)
        topic_names = extract_teaching_plan_knows(arrangementID=arrangementID)

        # 获取学生知识点集合
        all_knows = build_all_knows(student_data["knowledge_points"], extra_knows=topic_names)
        top5 = get_recommendations_and_log(student_data, all_knows)
        
        if top5:
            return jsonify({"status": "success", "recommendations": top5})
        return jsonify({"status": "error", "message": "没有推荐的题目"}), 400
    except Exception as e:
        return jsonify({"status": "error", "message": str(e)}), 500

# 具体知识点中调用的方法
@app.route("/recommend_weak_questions/teaching/knowledge/<topicID>", methods=["POST"])
def recommend_weak_questions_in_teaching_plan_with_topicID_endpoint(topicID):
    try:
        student_data = request.json
        student_data = map_knowledge_points(student_data, kp2id)
        topic_names = extract_teaching_plan_knows(topicID=topicID)
        # print("topic_names",topic_names)
        # 获取学生知识点集合
        all_knows = build_all_knows(student_data["knowledge_points"], extra_knows=topic_names)
        top5 = get_recommendations_and_log(student_data, all_knows)
        
        if top5:
            return jsonify({"status": "success", "recommendations": top5})
        return jsonify({"status": "error", "message": "没有推荐的题目"}), 400
    except Exception as e:
        return jsonify({"status": "error", "message": str(e)}), 500


# 教学计划完成后调用的方法
@app.route("/recommend_weak_questions/teaching/<programmeID>", methods=["POST"])
def recommend_weak_questions_after_teaching_plan_endpoint(programmeID):
    try:
        student_data = request.json
        student_data = map_knowledge_points(student_data, kp2id)
        topic_names = extract_teaching_plan_knows(programmeID=programmeID)

        # 获取学生知识点集合
        all_knows = build_all_knows(student_data["knowledge_points"], extra_knows=topic_names)
        top5 = get_recommendations_and_log(student_data, all_knows)
        
        if top5:
            return jsonify({"status": "success", "recommendations": top5})
        return jsonify({"status": "error", "message": "没有推荐的题目"}), 400
    except Exception as e:
        return jsonify({"status": "error", "message": str(e)}), 500
    
# 教学计划调用的方法 - 预习
@app.route("/recommend_weak_questions/teaching/pre/<programmeID>", methods=["POST"])
def recommend_weak_questions_pre_teaching_plan_endpoint(programmeID):
    try:
        student_data = request.json
        student_data = map_knowledge_points(student_data, kp2id)
        topic_names = extract_teaching_plan_knows(programmeID=programmeID)

        # 获取学生知识点集合
        # all_knows = build_all_knows(student_data["knowledge_points"], extra_knows=topic_names)
        all_knows = topic_names
        
        top5 = get_recommendations_and_log(student_data, all_knows, True)
        
        if top5:
            return jsonify({"status": "success", "recommendations": top5})
        return jsonify({"status": "error", "message": "没有推荐的题目"}), 400
    except Exception as e:
        return jsonify({"status": "error", "message": str(e)}), 500
    
# 提示词调用方法
@app.route("/softprompt", methods=["POST"])
def generate_prompt_endpoint():
    try:
        student_data = request.json
        parsed_data = parse_server_request(student_data)
        
        max_step = 50
        num_questions = 157
        
        soft_prompt = generate_final_prompt(
            parsed_data, dkt_model, knows_tree_path, kp_map_path, question_bank_csv_path, max_step, num_questions
        )
        

        if soft_prompt:
            return jsonify({"status": "success", "softprompt": soft_prompt})
        return jsonify({"status": "error", "message": "没有得到提示词"}), 400

    except Exception as e:
        return jsonify({"status": "error", "message": str(e)}), 500

@app.route("/")
def hello():
    return "Welcome to the DKT Recommendation Service!"

if __name__ == "__main__":
    app.run(debug=True, host="0.0.0.0", port=5000)
